{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "f27fcffd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import display\n",
    "import torch\n",
    "import os\n",
    "import pandas as pd\n",
    "from torch_geometric.data import InMemoryDataset\n",
    "from urllib.request import urlopen\n",
    "from torch_geometric.utils import negative_sampling\n",
    "import numpy as np\n",
    "from torch_geometric.nn import GCNConv\n",
    "from sklearn.metrics import roc_auc_score\n",
    "import torch_geometric.transforms as T\n",
    "from torch_geometric.transforms import ToUndirected, RandomLinkSplit\n",
    "from torch_geometric.data import Data\n",
    "from tqdm.notebook import tqdm\n",
    "from torch_geometric.loader import DataLoader\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "11d9cd9f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define Neo4j connections\n",
    "import pandas as pd\n",
    "from neo4j import GraphDatabase\n",
    "host = 'bolt://188.230.149.51:7687'\n",
    "user = 'neo4j'\n",
    "password = 'letmein'\n",
    "driver = GraphDatabase.driver(host,auth=(user, password))\n",
    "\n",
    "def run_query(query):\n",
    "    with driver.session() as session:\n",
    "        result = session.run(query)\n",
    "        return pd.DataFrame([r.values() for r in result], columns=result.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0498a537",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_frames():\n",
    "    \"\"\"\n",
    "    Loads the nodes and edges from Neo4j.\n",
    "    :return: nodes and edges frame\n",
    "    \"\"\"\n",
    "    dfn = run_query(\"\"\"\n",
    "    MATCH (u:User)\n",
    "    RETURN u.id AS id, u.age AS age, u.gender AS gender\n",
    "    \"\"\")\n",
    "    dfn = dfn.set_index(\"id\")\n",
    "\n",
    "    dfe = run_query(\"\"\"\n",
    "    MATCH (s:User)-[:FRIEND]->(t:User)\n",
    "    RETURN s.id as source_id, t.id as target_id\n",
    "    \"\"\")\n",
    "\n",
    "    return dfn, dfe\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4ee78325",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>source_id</th>\n",
       "      <th>target_id</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>16</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>10</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1</td>\n",
       "      <td>12</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1</td>\n",
       "      <td>8</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1</td>\n",
       "      <td>7</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  source_id target_id\n",
       "0         1        16\n",
       "1         1        10\n",
       "2         1        12\n",
       "3         1         8\n",
       "4         1         7"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>age</th>\n",
       "      <th>gender</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>26</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16</th>\n",
       "      <td>23</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>29</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>26</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>17</th>\n",
       "      <td>27</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    age  gender\n",
       "id             \n",
       "1    26       1\n",
       "16   23       1\n",
       "3    29       1\n",
       "4    26       0\n",
       "17   27       0"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "dfn, dfe = load_frames()\n",
    "display(dfe.head())\n",
    "display(dfn.head())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d46fa5af",
   "metadata": {},
   "outputs": [],
   "source": [
    "class IdentityEncoder(object):\n",
    "    # The 'IdentityEncoder' takes the raw column values and converts them to\n",
    "    # PyTorch tensors.\n",
    "    def __init__(self, dtype=None):\n",
    "        self.dtype = dtype\n",
    "\n",
    "    def __call__(self, df):\n",
    "        return torch.from_numpy(df.values).view(-1, 1).to(self.dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e5ec1425",
   "metadata": {},
   "outputs": [],
   "source": [
    "def transform_nodes(df, encoders = None, **kwargs):\n",
    "    \"\"\"\n",
    "    Transforms the node frame to a Torch payload and a mapping from the original frame index.\n",
    "    :param df: the nodes frame\n",
    "    :param encoders: encoders to be applied to the columns\n",
    "    :param kwargs: additional stuff\n",
    "    :return:\n",
    "    \"\"\"\n",
    "    # required mapping\n",
    "    mapping = {index: i for i, index in enumerate(df.index.unique())}\n",
    "    x = None\n",
    "    # apply encoding if specified\n",
    "    if encoders is not None:\n",
    "        xs = [encoder(df[col]) for col, encoder in encoders.items()]\n",
    "        x = torch.cat(xs, dim = -1)\n",
    "    return x, mapping"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "ba1aeb23",
   "metadata": {},
   "outputs": [],
   "source": [
    "nodes_x, nodes_mapping = transform_nodes(dfn, \n",
    "        encoders = {\"age\": IdentityEncoder(torch.float), \"gender\": IdentityEncoder(torch.float)})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "5b5a603f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def transform_edges(df_edges, nodes_mapping, encoders = None, **kwargs):\n",
    "    \"\"\"\n",
    "    Transforms the edge frame to a Pyg compatible index and payload.\n",
    "    :param df_edges: the edge frame\n",
    "    :param nodes_mapping: the node index map\n",
    "    :param encoders: optional encoders\n",
    "    :param kwargs:\n",
    "    :return:\n",
    "    \"\"\"\n",
    "    src = [nodes_mapping[src_id] for src_id in df_edges.source_id]\n",
    "    dst = [nodes_mapping[tgt_id] for tgt_id in df_edges.target_id]\n",
    "    edge_index = torch.tensor([src, dst])\n",
    "\n",
    "    edge_attr = None\n",
    "    if encoders is not None:\n",
    "        edge_attrs = [encoder(df_edges[col]) for col, encoder in encoders.items()]\n",
    "        edge_attr = torch.cat(edge_attrs, dim = -1)\n",
    "\n",
    "    return edge_index, edge_attr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "3c765367",
   "metadata": {},
   "outputs": [],
   "source": [
    "edges_index, edges_label = transform_edges(\n",
    "    dfe,\n",
    "    nodes_mapping = nodes_mapping\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "e9adab63",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_data(nodes_payload, edges_index):\n",
    "    d = Data(x = nodes_payload, edge_index = edges_index)\n",
    "    return d\n",
    "\n",
    "data = create_data(nodes_x, edges_index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "01f35af4",
   "metadata": {},
   "outputs": [],
   "source": [
    "transform = T.Compose([\n",
    "    T.NormalizeFeatures(),\n",
    "    T.ToDevice(device),\n",
    "    T.RandomLinkSplit(num_val = 0.05, num_test = 0.1, is_undirected = True,\n",
    "                      add_negative_train_samples = False),\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "ccb30476",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data, val_data, test_data = transform(data)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "6451c125",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Net(torch.nn.Module):\n",
    "    def __init__(self, in_channels, hidden_channels, out_channels):\n",
    "        super().__init__()\n",
    "        self.conv1 = GCNConv(in_channels, hidden_channels)\n",
    "        self.conv2 = GCNConv(hidden_channels, out_channels)\n",
    "\n",
    "    def encode(self, x, edge_index):\n",
    "        # chaining two convolutions with a standard relu activation\n",
    "        x = self.conv1(x, edge_index).relu()\n",
    "        return self.conv2(x, edge_index)\n",
    "\n",
    "    def decode(self, z, edge_label_index):\n",
    "        # cosine similarity\n",
    "        return (z[edge_label_index[0]] * z[edge_label_index[1]]).sum(dim = -1)\n",
    "\n",
    "    def decode_all(self, z):\n",
    "        # decode the whole lot in one go\n",
    "        prob_adj = z @ z.t()\n",
    "        return (prob_adj > 0).nonzero(as_tuple = False).t()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "570d4707",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Net(data.num_features, 128, 64).to(device)\n",
    "optimizer = torch.optim.Adam(params = model.parameters(), lr = 0.01)\n",
    "criterion = torch.nn.BCEWithLogitsLoss()\n",
    "\n",
    "\n",
    "def train():\n",
    "    model.train()\n",
    "    optimizer.zero_grad()\n",
    "    z = model.encode(train_data.x, train_data.edge_index)\n",
    "\n",
    "    # as explained, some negative sampling is used in each epoch\n",
    "    neg_edge_index = negative_sampling(edge_index = train_data.edge_index, num_nodes = train_data.num_nodes,\n",
    "                                       num_neg_samples = train_data.edge_label_index.size(1), method = 'sparse')\n",
    "\n",
    "    edge_label_index = torch.cat([train_data.edge_label_index, neg_edge_index],\n",
    "                                 dim = -1,\n",
    "                                 )\n",
    "    edge_label = torch.cat([train_data.edge_label, train_data.edge_label.new_zeros(neg_edge_index.size(1))\n",
    "                            ], dim = 0)\n",
    "\n",
    "    out = model.decode(z, edge_label_index).view(-1)\n",
    "    loss = criterion(out, edge_label)\n",
    "    # standard torch mechanics here\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "1736c38a",
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def test(d):\n",
    "    model.eval()\n",
    "    z = model.encode(d.x, d.edge_index)\n",
    "    out = model.decode(z, d.edge_label_index).view(-1).sigmoid()\n",
    "    return roc_auc_score(d.edge_label.cpu().numpy(), out.cpu().numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f067d6e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "final_test_auc = 0\n",
    "best_val_auc = 0\n",
    "epochs = 3\n",
    "for epoch in range(1, epochs + 1):\n",
    "    loss = train()\n",
    "    # validation accuracy\n",
    "    val_auc = test(val_data)\n",
    "    # test accuracy\n",
    "    test_auc = test(test_data)\n",
    "    if val_auc > best_val_auc:\n",
    "        best_val = val_auc\n",
    "        final_test_auc = test_auc\n",
    "    print(f'Epoch: {epoch:03d}, Loss: {loss:.4f}, Val: {val_auc:.4f}, '\n",
    "          f'Test: {test_auc:.4f}')\n",
    "\n",
    "print(f'Final Test: {final_test_auc:.4f}')\n",
    "# the final test accuracy\n",
    "z = model.encode(test_data.x, test_data.edge_index)\n",
    "final_edge_index = model.decode_all(z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a38ed30",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
